#include "OptimusMaterius.h"
#include "Ray.h"
#include "Scene.h"

OptimusMaterius::OptimusMaterius(const Vector3 & kd, const Vector3 & ks, const float shinyness, const Vector3 & kt, float (*rIndex)(short), const Vector3 & ka) :
    m_kd(kd), m_ks(ks), m_shinyness(shinyness), m_kt(kt), refIndex(rIndex), m_ka(ka)
{
}

OptimusMaterius::~OptimusMaterius()
{
}

Vector3 OptimusMaterius::shade(const Ray& ray, const HitInfo& hit, char rayCount, const Scene& scene, int mode, Vector3 rayColor, bool diffuseRay) const
{
    Vector3 L = Vector3(0.0f, 0.0f, 0.0f);
    
	const Vector3 viewDir = -ray.d; // d is already normalized
	Vector3 reflection = ray.d - 2*(dot(ray.d,hit.N))*hit.N;
	reflection.normalize();

    ////////////// create measurement point
    if (mode == APPM_MEAS && m_kd.length2() > epsilon && rayCount < MAX_REFLECTIONS)
    {
    	HPoint *hp = new HPoint;
    	hp->isLight = false;
    	hp->f = m_kd * rayColor;
    	hp->pos = hit.P;
    	hp->norm = hit.N;
    	hp->pix = g_pixel_index;
    	g_hitpoints = g_hitpoints->ListAdd(hp);
    }

	if (mode != APPM_MEAS)
	{
		//////////////// compute direct diffuse componant of lighting

		// compute lighting from the point lights
		const Lights *lightlist = scene.lights();
		Lights::const_iterator lightIter;
		PointLight* pLight;
		for (lightIter = lightlist->begin(); lightIter != lightlist->end(); lightIter++)
		{
			pLight = *lightIter;

			Vector3 l = pLight->position() - hit.P;

			// the inverse-squared falloff
			float falloff = l.length2();

			// normalize the light direction
			l /= sqrt(falloff);

			// get the diffuse and highlight componants
			float nDotL = dot(hit.N, l);
			Vector3 result = pLight->color();

			//Check if there should be a shadow
			Ray lightRay(hit.P, l.normalized());
			HitInfo newHitLight;
			if (!scene.trace(newHitLight, lightRay, epsilon)) {
				// this is +1 if reflection points right at light
				// don't let it go negative for when it points away
				Vector3 shine(0);
				if (m_shinyness > epsilon)
				{
					float reflect_towards_light = std::max(0.0f,dot(l, reflection));
					shine = m_ks*(pow(reflect_towards_light, m_shinyness));
				}
				result *= m_kd + shine;
				L += std::max(0.0f, (pLight->wattage()*nDotL/(4*PI*falloff))) * result;
			}
		}

		// compute lighting from the area lights
		const AreaLights *alightlist = scene.areaLights();
		AreaLights::const_iterator alightIter;
		AreaLight* pALight;
		for (alightIter = alightlist->begin(); alightIter != alightlist->end(); alightIter++)
		{
			pALight = *alightIter;
			L += m_kd * pALight->sampleLight(hit, scene);
		}

		
		/////////////////// add the ambient component

		L += m_ka;
	}

	/////////////// compute indirect diffuse lighting

	if (mode == PATHTRACE && rayCount < MAX_REFLECTIONS)
	{
		// trace a random path
		Ray newRay(hit.P, Material::randDir(hit.N, ray));
		HitInfo newHitInfo;
		if (scene.trace(newHitInfo, newRay, epsilon))
		{
			L += m_kd * newHitInfo.material->shade(newRay, newHitInfo, rayCount+1, scene, mode, rayColor, true);
		}
	}
    

	/////////////// compute specular componant of lighting (reflection and transmission)

	if(rayCount < MAX_REFLECTIONS)
	{
		float Fr=0, Ft=0;
		if(refIndex)
		{			
            //indexes of refraction
            float n1, n2;

			bool isInside = !(dot(hit.N,ray.d)<epsilon);
            Vector3 norm = hit.N;
            Vector3 transparency;
            // are we going from air to material
            if(!isInside)
			{
                n1 = IOR_AIR;
				n2 = refIndex(-1);
				transparency = m_kt;
            }
            //Otherwise going from material to air
            else 
			{
				n1 = refIndex(-1);
                n2 = IOR_AIR;
                norm = -1.0f*hit.N;
                transparency = Vector3(1);
            }

			float n3 = n1/n2;
			Vector3 w = -1.0*ray.d;
			float dotProd = dot(w,norm);
		
			float innerRoot = 1.0f-((n3)*(n3))*(1.0f-(dotProd*dotProd));

			Vector3 refraction;
			if(innerRoot>=0) {
				refraction = -1.0f*(n3)*(w - dotProd*norm)-sqrt(innerRoot)*norm;
			}
			else {
				refraction = reflection;
			}


			// Fresnel

			float cosTheta1 = fabs(dotProd); //Calculate cos(theta1)
			//float sinTheta1Sqrd = 1.0f - cosTheta1*cosTheta1; //Calculate sin^2(theta1)
			//float cosTheta2 = sqrt(1.0f - n3*n3*sinTheta1Sqrd); //Calculate cos(theta2), should be equivalent to above 3 lines, but way more efficient
			//float p_parallel = (n2*cosTheta1 - n1*cosTheta2)/(n2*cosTheta1 + n1*cosTheta2);
			//float p_perpendicular = (n1*cosTheta1 - n2*cosTheta2)/(n1*cosTheta1 + n2*cosTheta2);
			//float Fr = 0.5f*(p_parallel*p_parallel + p_perpendicular*p_perpendicular); //Final calculation of Fresnel equation

			// approximation - probably good enough, and a little faster to compute
			float Fo = ((n1 - n2)*(n1 - n2)) / ((n1 + n2)*(n1 + n2));
			Fr = Fo + (1 - Fo)*pow((1 - cosTheta1), 5);
			Ft = 1 - Fr;

			Vector3 L2 = L;
			refraction.normalize();
			Ray newRay2(hit.P, refraction);
			HitInfo newHitInfo2;


			if (scene.trace(newHitInfo2, newRay2, epsilon))
			{
				L += Ft * newHitInfo2.material->shade(newRay2, newHitInfo2, (rayCount+1), scene, mode, Ft * rayColor);
			}
			else
			{
				// environment map it
				refraction.normalize();
				L += Ft * scene.getMapValue(refraction);
			}
		}
		
		Ray newRay(hit.P, reflection);
		HitInfo newHitInfo;
		if (scene.trace(newHitInfo, newRay, epsilon))
		{
			L += (m_ks + m_kt*Fr) * newHitInfo.material->shade(newRay, newHitInfo, rayCount+1, scene, mode, (m_ks + m_kt*Fr) * rayColor);
		}
		else
		{
			// environment map it
			reflection.normalize();
			L += (m_ks + m_kt*Fr) * scene.getMapValue(reflection);
		}

	}

    return L;

}

void OptimusMaterius::handlePhoton(const Ray& ray, const HitInfo& hit, char rayCount, Vector3 photonFlux, const Scene& scene, int mode, LightPath *lp) const
{
    if (mode == APPM_PHOTON && rayCount < MAX_REFLECTIONS)
    {
    	// diffuse
    	if (m_kd.length() > epsilon)
    	{
        	// from Toshiya's PPM code
    		Vector3 hh = (hit.P - g_hpbbox.min) * g_hash_s;
    		int ix = abs(int(hh.x)), iy = abs(int(hh.y)), iz = abs(int(hh.z));
    		// strictly speaking, we should use #pragma omp critical here
    		{
    			List* hp = g_hash_grid[scene.hash(ix, iy, iz)];
    			while (hp != NULL)
    			{
    				HPoint *measPt = hp->id;
    				hp = hp->next;
    				Vector3 v = measPt->pos - hit.P;

    				if (dot(ray.d, measPt->norm) < -epsilon && (dot(measPt->norm, hit.N) > epsilon) &&
    						(v.length2() <= measPt->r2))
    				{
    					double g = (measPt->n * ALPHA + ALPHA) / (measPt->n * ALPHA + 1.0);
    					measPt->r2 = measPt->r2 * g;
    					measPt->n++;
    					measPt->flux = (measPt->flux + measPt->f * photonFlux * (1. / PI)) * g;
    				}
    			}
    		}
    		Vector3 f = m_kd;
        	double p = f.x > f.y && f.x > f.z ? f.x : f.y > f.z ? f.y : f.z;
    		if (dot(ray.d, hit.N) < -epsilon && Rand::getRand() < p) // russian roullette
    		{
    			// then trace a random path
    			//Ray newRay(hit.P, Material::randDir(hit.N, ray));
    			float u1, u2;
    			if (lp)
    				lp->getDirVars(rayCount, u1, u2);
    			else
    			{
    				u1 = Rand::getRand();
    				u2 = Rand::getRand();
    			}
    			Ray newRay(hit.P, Rand::getRandDir(hit.N, u1, u2));
    			HitInfo newHitInfo;
    			if (scene.trace(newHitInfo, newRay, epsilon))
    			{
    				newHitInfo.material->handlePhoton(newRay, newHitInfo, rayCount+1, m_kd * photonFlux * (1./p), scene, mode, lp);
    			}
    		}
    	}


		float Fr=0, Ft=0;
		Vector3 reflection = ray.d - 2*(dot(ray.d,hit.N))*hit.N;
		reflection.normalize();

		// transmission
		if(refIndex)
		{
            //indexes of refraction
            float n1, n2;

			bool isInside = !(dot(hit.N,ray.d)<epsilon);
            Vector3 norm = hit.N;
            Vector3 transparency;
            // are we going from air to material
            if(!isInside)
			{
                n1 = IOR_AIR;
				n2 = refIndex(-1);
				transparency = m_kt;
            }
            //Otherwise going from material to air
            else
			{
				n1 = refIndex(-1);
                n2 = IOR_AIR;
                norm = -1.0f*hit.N;
                transparency = Vector3(1);
            }

			float n3 = n1/n2;
			Vector3 w = -1.0*ray.d;
			float dotProd = dot(w,norm);

			float innerRoot = 1.0f-((n3)*(n3))*(1.0f-(dotProd*dotProd));

			Vector3 refraction;
			if(innerRoot>=0) {
				refraction = -1.0f*(n3)*(w - dotProd*norm)-sqrt(innerRoot)*norm;
			}
			else {
				refraction = reflection;
			}


			// Fresnel

			float cosTheta1 = fabs(dotProd);

			// approximation - probably good enough, and a little faster to compute
			float Fo = ((n1 - n2)*(n1 - n2)) / ((n1 + n2)*(n1 + n2));
			Fr = Fo + (1 - Fo)*pow((1 - cosTheta1), 5);
			Ft = 1 - Fr;

			refraction.normalize();
			Ray newRay(hit.P, refraction);
			HitInfo newHitInfo;


			if (scene.trace(newHitInfo, newRay, epsilon))
			{
				newHitInfo.material->handlePhoton(newRay, newHitInfo, rayCount+1, Ft * photonFlux, scene, mode, lp);
			}
		}

		// reflection
    	if (m_ks.length() > epsilon)
    	{

    		Ray newRay(hit.P, reflection);
    		HitInfo newHitInfo;
    		if (scene.trace(newHitInfo, newRay, epsilon))
    		{
    			newHitInfo.material->handlePhoton(newRay, newHitInfo, rayCount+1, (m_ks + m_kt*Fr) * photonFlux, scene, mode, lp);
    		}
    	}
    }
}

